import torch
import torch.nn as nn
from torchvision import models

#此为PyTorch的神经网络模型，利用预训练的 ResNet-34 提取特征，并通过转置卷积进行上采样，从而生成语义分割的预测结果
class FCN(nn.Module):
    def __init__(self, num_classes):
        super().__init__()
        # 获取未经预训练的基础网络 ResNet-34 模型(34层的残差网络)
        pretrained_net = models.resnet34(weights = None) # False表示不加载预训练模型的参数

        # 获取三个池化层，通过提取不同阶段的特征图获得不同尺度的特征信息
        # 提取输入图像的低级特征，排除最后四个池化层
        self.stage1 = nn.Sequential(*list(pretrained_net.children())[:-4])
        # 提取输入图像的中级特征，获取倒数第四个模块
        self.stage2 = list(pretrained_net.children())[-4]
        # 提取输入图像的高级特征，获取倒数第三个模块
        self.stage3 = list(pretrained_net.children())[-3]

        # 获得单通道的预测输出，计算三个尺度的分数
        # 此卷积层期望输入特征图128个通道，输出特征图num_classes个通道，卷积核为1*1
        self.scores1 = nn.Conv2d(128, num_classes, 1)
        # 此卷积层期望输入特征图256个通道，输出特征图num_classes个通道，卷积核为1*1
        self.scores2 = nn.Conv2d(256, num_classes, 1)
        # 此卷积层期望输入特征图512个通道，输出特征图num_classes个通道，卷积核为1*1
        self.scores3 = nn.Conv2d(512, num_classes, 1)

        # 8倍上采样，输入和输出通道数都为num_classes，卷积核大小为16，步长为8，填充为4
        self.upsample_8x = nn.ConvTranspose2d(num_classes, num_classes, 16, 8, 4, bias=False)
        # 2倍上采样，输入和输出通道数都为num_classes，卷积核大小为4，步长为2，填充为1
        self.upsample_2x = nn.ConvTranspose2d(num_classes, num_classes, 4, 2, 1, bias=False)
        # sigmoid激活函数用于最终的输出结果，将特征图中的像素值压缩到[0,1]，用于表示每个像素属于目标的概率
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        x = self.stage1(x)
        s1 = x  # 获取1/8池化输出

        x = self.stage2(x)
        s2 = x  # 获取1/16池化输出

        x = self.stage3(x)
        s3 = x  # 获取1/32池化输出

        # 计算1/32分数
        s3_scores = self.scores3(s3)
        # 上采样2倍，并进行融合
        s3_scores_x2 = self.upsample_2x(s3_scores)
        s2_scores = self.scores2(s2)
        s2_fuse = s2_scores + s3_scores_x2

        # 计算1/8分数
        s1_scores = self.scores1(s1)
        # 上采样2倍，并进行融合
        s2_fuse_x2 = self.upsample_2x(s2_fuse)
        s = s1_scores + s2_fuse_x2

        # 上采样8倍，获取的s_x8与原始输入尺寸相同
        s_x8 = self.upsample_8x(s)
        s_x8 = self.sigmoid(s_x8)
        return s_x8


if __name__ == "__main__":
    fcn = FCN(num_classes=1) # 实例化FCN模型
    x = torch.randn(size=(1, 3, 224, 224)) # 随机生成输入张量
    out = fcn(x) # 输入张量，得到输出张量
    print(out.size()) # 打印输出张量的尺寸